/*******************************************************************************
 * Copyright (c) 2014 Steffen Mikkelsen
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the GNU Public License v3.0
 * which accompanies this distribution, and is available at
 * http://www.gnu.org/licenses/gpl.html
 * 
 * Contributors:
 *     Steffen Mikkelsen - initial API and implementation
 ******************************************************************************/
package com.subsecret.secureconnection.securestream.encryption;

import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.KeyAgreement;
import javax.crypto.interfaces.DHPrivateKey;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;

/**
 * Diffie Hellman implementation for Server-Client communication
 * 
 * @author steffen
 * 
 */
public class DiffieHellmanClient {
    
    private byte[] secret = null;
    private DHPrivateKey receiverPrivate = null;
    
    public DiffieHellmanClient() {
    }
    
    /**
     * Process client part of Diffie Hellman.
     * 
     * @param receivedPublic
     *            public key from server
     * @return public key to be sent back to server
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     * @throws InvalidAlgorithmParameterException
     * @throws InvalidKeyException
     */
    public byte[] generateClientPart(byte[] receivedPublic) {
        X509EncodedKeySpec x509Key = new X509EncodedKeySpec(receivedPublic);
        KeyFactory kfactory;
        try {
            kfactory = KeyFactory.getInstance("DiffieHellman");
            DHPublicKey receivedPublicKey = (DHPublicKey) kfactory.generatePublic(x509Key);
            DHParameterSpec receivedDHParams = receivedPublicKey.getParams();
            
            KeyPairGenerator kpg = KeyPairGenerator.getInstance("DiffieHellman");
            kpg.initialize(receivedDHParams);
            KeyPair receiverKeyPair = kpg.generateKeyPair();
            DHPublicKey receiverPublic = (DHPublicKey) receiverKeyPair.getPublic();
            receiverPrivate = (DHPrivateKey) receiverKeyPair.getPrivate();
            
            // share public key with Server
            byte[] encodedPublic = receiverPublic.getEncoded();
            
            KeyAgreement senderAgree = KeyAgreement.getInstance("DiffieHellman");
            senderAgree.init(receiverPrivate);
            senderAgree.doPhase(receivedPublicKey, true);
            secret = senderAgree.generateSecret();
            return encodedPublic;
        }
        catch (NoSuchAlgorithmException | InvalidKeyException | IllegalStateException | InvalidAlgorithmParameterException | InvalidKeySpecException e) {
            e.printStackTrace();
        }
        return null;
    }
    
    /**
     * @return The shared secret, available after generateClientPart has been called
     */
    public byte[] getSharedSecret() {
        return secret;
    }
    
}
